import concurrent.futures
import threading
import pathlib
import io
import PIL
import cv2
import numpy as np
from semantic_nav.log_utils import get_logger
logger = get_logger()


def image_to_bytes(img_arg):
    if isinstance(img_arg, str) or isinstance(img_arg, pathlib.Path):
        # Open the image and convert it to a byte stream
        with open(img_arg, "rb") as img:
            img_byte = img.read()
    elif isinstance(img_arg, PIL.Image.Image):
        # Create a byte stream
        buffered = io.BytesIO()
        # Save the image to a byte stream in PNG format
        img_arg.save(buffered, format="PNG")
        # Gets byte data from a byte stream
        img_byte = buffered.getvalue()
    elif isinstance(img_arg, np.ndarray):
        _, buffer = cv2.imencode('.png', img_arg)
        img_byte = buffer.tobytes()
    elif isinstance(img_arg, bytes):
        img_byte = img_arg
    else:
        raise ValueError("Unsupported type for image message") 
    return img_byte


class WrappedComparable:
    def __init__(self, priority, value):
        self.priority = priority
        self.value = value

    def __lt__(self, other):
        return self.priority < other.priority

    def __eq__(self, other):
        return self.priority == other.priority

    def __str__(self):
        return f"WrappedComparable(value={self.value}, priority={self.priority})"


class SingleTaskExecutor:
    def __init__(self):
        self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
        self.current_future = None
        self.wait_task = None
        self.latest_result = None
        self.latest_exception = None
        self.result_ready = threading.Event()
        self.lock = threading.RLock()

    def submit(self, fn, *args, **kwargs):
        with self.lock:
            if self.current_future is None or self.current_future.done():
                self.result_ready.clear()
                self.current_future = self.executor.submit(self._run_task, fn, *args, **kwargs)
            else:
                self.wait_task = (fn, args, kwargs)
        return self.current_future

    def _run_task(self, fn, *args, **kwargs):
        try:
            result = fn(*args, **kwargs)
            with self.lock:
                logger.info(f"Task completed successfully. Result: {result}")
                self.latest_result = result
                self.latest_exception = None
            return result
        except Exception as e:
            logger.error(f"Task failed with exception: {e}")
            with self.lock:
                self.latest_result = None
                self.latest_exception = e
            raise
        finally:
            self.result_ready.set()
            with self.lock:
                if self.wait_task:
                    next_task = self.wait_task
                    self.wait_task = None
                    self.submit(*next_task)

    def result(self, timeout=None):
        if not self.result_ready.wait(timeout=timeout):
            raise concurrent.futures.TimeoutError("Task execution timed out")

        with self.lock:
            if self.latest_exception:
                raise self.latest_exception
            return self.latest_result

    def cancel(self):
        with self.lock:
            if self.current_future:
                self.current_future.cancel()
            self.wait_task = None

    def shutdown(self):
        self.cancel()
        self.executor.shutdown()
